package utils

import (
	"bufio"
	"bytes"
	"context"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/hex"
	"encoding/xml"
	"errors"
	"fmt"
	"github.com/redis/go-redis/v9"
	"github.com/zeromicro/go-zero/core/logx"
	"github.com/zeromicro/go-zero/core/stores/sqlx"
	"github.com/zeromicro/go-zero/zrpc"
	"golang.org/x/crypto/ssh"
	"io"
	"net"
	"os"
	"os/exec"
	"strings"
	"sync"
	"time"
	"yunzhan/common/models"
)

/**
* @Author: way lee
* @File: utils.go
* @Description:
* @Date: 2024/6/24 17:28
 */

// SSHConnectionPool 维护一个 SSH 连接池，允许复用连接以提高性能。
type SSHConnectionPool struct {
	connections chan *ssh.Client  // 用于存储 SSH 连接的通道
	mu          sync.Mutex        // 保护连接池的互斥锁
	config      *ssh.ClientConfig // SSH 客户端配置
	host        string            // SSH 服务器主机名
	port        string            // SSH 服务器端口
}

// SafeNewClient 用于安全地创建 zrpc.Client
func SafeNewClient(host, port string) (zrpc.Client, error) {
	defer func() {
		if r := recover(); r != nil {
			logx.Errorf("创建RPC客户端失败: %v", r)
		}
	}()
	return zrpc.MustNewClient(zrpc.RpcClientConf{
		Target:  fmt.Sprintf("dns:///%s:%s", host, port),
		Timeout: int64((30 * time.Second) / time.Millisecond),
	}), nil
}

// NewSSHConnectionPool 创建一个新的 SSH 连接池
func NewSSHConnectionPool(host, port, user, password string, maxConnections int) *SSHConnectionPool {
	config := &ssh.ClientConfig{
		User:            user,
		Auth:            []ssh.AuthMethod{ssh.Password(password)},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         5 * time.Second,
	}

	pool := &SSHConnectionPool{
		connections: make(chan *ssh.Client, maxConnections),
		config:      config,
		host:        host,
		port:        port,
	}

	for i := 0; i < maxConnections; i++ {
		client, err := ssh.Dial("tcp", host+":"+port, config)
		if err != nil {
			logx.Errorf("创建 SSH 连接失败: %v", err)
			continue
		}
		pool.connections <- client
	}

	return pool
}

// GetConnection 从连接池中获取一个可用的 SSH 连接
func (p *SSHConnectionPool) GetConnection() (*ssh.Client, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	select {
	case client := <-p.connections:
		// 检查连接是否有效
		_, _, err := client.SendRequest("keepalive@openssh.com", true, nil)
		if err != nil {
			client.Close()
			newClient, err := ssh.Dial("tcp", p.host+":"+p.port, p.config)
			if err != nil {
				return nil, fmt.Errorf("创建 SSH 连接失败: %v", err)
			}
			return newClient, nil
		}
		return client, nil
	case <-ctx.Done():
		return nil, errors.New("timeout: could not get SSH connection from pool")
	}
}

// ReleaseConnection 释放一个 SSH 连接回连接池
func (p *SSHConnectionPool) ReleaseConnection(client *ssh.Client) {
	p.mu.Lock()
	defer p.mu.Unlock()
	select {
	case p.connections <- client:
	default:
		// 如果池已满，关闭连接
		client.Close()
	}
}

// Close 关闭连接池中的所有连接
func (p *SSHConnectionPool) Close() {
	p.mu.Lock()
	defer p.mu.Unlock()

	var wg sync.WaitGroup
	close(p.connections)
	for client := range p.connections {
		wg.Add(1)
		go func(c *ssh.Client) {
			defer wg.Done()
			c.Close()
		}(client)
	}
	wg.Wait()
}

func ExecCommand(command string) (string, error) {
	cmd := exec.Command("bash", "-c", command)
	cmd.Env = append(os.Environ(), "LC_ALL=C", "LANG=C")

	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	err := cmd.Run()
	if err != nil {
		errMsg := fmt.Sprintf("命令执行失败: %s, stdout: %s, stderr: %s, 错误: %v", command, stdout.String(), stderr.String(), err)
		return "", fmt.Errorf(errMsg)
	}

	return stdout.String(), nil
}

// Rtrim 去除字符串右侧的指定字符
func Rtrim(s string, flag string) string {
	s = strings.TrimSpace(s)
	if strings.HasSuffix(s, flag) {
		s = strings.TrimRight(s, flag)
	}
	return s
}

// GetCurrentInternalIP 获取当前机器的内网 IP 地址
func GetCurrentInternalIP() (string, error) {
	interfaces, err := net.Interfaces()
	if err != nil {
		return "", err
	}

	for _, iface := range interfaces {
		if iface.Flags&net.FlagUp == 0 || iface.Flags&net.FlagLoopback != 0 {
			continue
		}

		adders, err := iface.Addrs()
		if err != nil {
			return "", err
		}

		for _, addr := range adders {
			ipNet, isValidIPNet := addr.(*net.IPNet)
			if isValidIPNet && !ipNet.IP.IsLoopback() && ipNet.IP.To4() != nil {
				return ipNet.IP.String(), nil
			}
		}
	}

	return "", errors.New("未找到有效的内网IP地址")
}

// CheckUbuntu 检查当前操作系统是否为 Ubuntu
func CheckUbuntu() (bool, error) {
	file, err := os.Open("/etc/os-release")
	if err != nil {
		return false, err
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		if strings.Contains(scanner.Text(), "Ubuntu") {
			return true, nil
		}
	}

	if err := scanner.Err(); err != nil {
		return false, err
	}

	return false, nil
}

// GetCurrentHostname 获取当前 Linux 主机名
func GetCurrentHostname() (string, error) {
	hostname, err := os.Hostname()
	if err != nil {
		return "", err
	}

	if hostname == "" {
		return "", errors.New("未找到有效的主机名")
	}

	return hostname, nil
}

// LogDeploymentStep 记录组件部署过程状态
func LogDeploymentStep(sqlConn sqlx.SqlConn, insertSql, taskID, nodeHost, componentName string, stepNumber int, stepDescription, status, errorMessage string) {
	logMessage := fmt.Sprintf("[Task ID: %s] [Node: %s] [Component: %s] [Step %d]: %s - %s. %s",
		taskID, nodeHost, componentName, stepNumber, stepDescription, status, errorMessage)

	if status == "failure" {
		logx.Errorf(logMessage)
	} else {
		logx.Infof(logMessage)
	}

	if _, err := sqlConn.Exec(insertSql, taskID, nodeHost, componentName, stepNumber, stepDescription, status, errorMessage); err != nil {
		logx.Errorf("%s部署日志记录失败: %v", componentName, err)
	}
}

// GetZooKeeperRole 获取 ZooKeeper 节点角色
func GetZooKeeperRole(host string) (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	conn, err := net.DialTimeout("tcp", host, 5*time.Second)
	if err != nil {
		return "", fmt.Errorf("连接 ZooKeeper 失败: %w", err)
	}
	defer conn.Close()

	_, err = conn.Write([]byte("srvr\n"))
	if err != nil {
		return "", fmt.Errorf("发送 srvr 命令失败: %w", err)
	}

	scanner := bufio.NewScanner(conn)
	var response string

	for scanner.Scan() {
		response += scanner.Text() + "\n"
		select {
		case <-ctx.Done():
			return "", ctx.Err()
		default:
			continue
		}
	}

	if err = scanner.Err(); err != nil {
		if errors.Is(err, os.ErrDeadlineExceeded) {
			return "", ctx.Err()
		}
		return "", fmt.Errorf("读取 ZooKeeper 响应失败: %w", err)
	}

	if strings.Contains(response, "Mode: leader") {
		return "Leader", nil
	} else if strings.Contains(response, "Mode: follower") {
		return "Follower", nil
	} else if strings.Contains(response, "Mode: observer") {
		return "Observer", nil
	}

	return "Unknown", nil
}

// LoadConfiguration 从文件加载 XML 配置
func LoadConfiguration(filename string) (*models.XmlConfiguration, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	var config models.XmlConfiguration
	if err = xml.NewDecoder(file).Decode(&config); err != nil {
		return nil, err
	}

	return &config, nil
}

// Contains 检查切片中是否包含指定字符串
func Contains(slice []string, str string) bool {
	for _, value := range slice {
		if value == str {
			return true
		}
	}
	return false
}

// Redis 连接池（单例模式）
var redisClient *redis.Client
var once sync.Once

// GetRedisClient 获取 Redis 客户端的单例实例
func GetRedisClient(redisAddr, redisPassword string) *redis.Client {
	once.Do(func() {
		redisClient = redis.NewClient(&redis.Options{
			Addr:     redisAddr,
			Password: redisPassword,
			DB:       0,
		})
	})
	return redisClient
}

// GetRedisRole 获取 Redis 角色信息
func GetRedisRole(redisAddr, redisPassword string) (string, error) {
	client := GetRedisClient(redisAddr, redisPassword)
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	result, err := client.Info(ctx, "replication").Result()
	if err != nil {
		return result, fmt.Errorf("获取 Redis 信息失败: %w", err)
	}

	lines := strings.Split(result, "\r\n")
	for _, line := range lines {
		if strings.HasPrefix(line, "role:") {
			role := strings.TrimSpace(strings.Split(line, ":")[1])
			return role, nil
		}
	}

	return "", fmt.Errorf("无法确定Redis角色")
}

// Encrypt 加密函数，使用 AES 算法
func Encrypt(key []byte, plaintext string) (string, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}

	ciphertext := make([]byte, aes.BlockSize+len(plaintext))
	iv := ciphertext[:aes.BlockSize]
	if _, err = io.ReadFull(rand.Reader, iv); err != nil {
		return "", err
	}

	stream := cipher.NewCFBEncrypter(block, iv)
	stream.XORKeyStream(ciphertext[aes.BlockSize:], []byte(plaintext))

	return hex.EncodeToString(ciphertext), nil
}

// Decrypt 解密函数，使用 AES 算法
func Decrypt(key []byte, encryptedText string) (string, error) {
	ciphertext, err := hex.DecodeString(encryptedText)
	if err != nil {
		return "", err
	}

	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}

	if len(ciphertext) < aes.BlockSize {
		return "", fmt.Errorf("密码太短(16)")
	}

	iv := ciphertext[:aes.BlockSize]
	ciphertext = ciphertext[aes.BlockSize:]

	stream := cipher.NewCFBDecrypter(block, iv)
	stream.XORKeyStream(ciphertext, ciphertext)

	return string(ciphertext), nil
}
